//*********************************************************
//
//    Copyright (c) Microsoft. All rights reserved.
//    This code is licensed under the Apache License, Version 2.0.
//    THIS CODE IS PROVIDED *AS IS* WITHOUT WARRANTY OF
//    ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY
//    IMPLIED WARRANTIES OF FITNESS FOR A PARTICULAR
//    PURPOSE, MERCHANTABILITY, OR NON-INFRINGEMENT.
//
//*********************************************************

namespace Trident.Scheduler.Hpc
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Text;
    using Microsoft.Research.eResearch.TridentScheduleCommon;
    using System.Reflection;
    using System.IO;
    using ER = Microsoft.Research.DataLayer;
    using System.Xml.Serialization;
    using System.Configuration;
    using System.Collections.ObjectModel;
    using Microsoft.Research.eResearch.Common.Logging;
    using Microsoft.Research.eResearch.Common.Eventing;
    using System.Globalization;
    using SH = Microsoft.Hpc.Scheduler;
    using SHP = Microsoft.Hpc.Scheduler.Properties;

    /// <summary>
    /// Trident HPC Scheduler class implements the ITridentScheduler interface 
    /// </summary>
    public class TridentHpcScheduler : ITridentScheduler, ITridentScheduleStatus
    {
        #region member data
        /// <summary>
        /// Member holding the current schedule engine
        /// </summary>
        private List<ITridentHpcScheduleEngine> currentScheduleEngines= null;
        /// <summary>
        /// Handle to registry database
        /// </summary>
        private ER.Connection registryConnection;
        /// <summary>
        /// Manager class for invoking HPC scheduler API
        /// </summary>
        private HPCSchedulerWrapper clusterManager = new HPCSchedulerWrapper();
        
        /// <summary>
        /// Workflow loader application name/Path
        /// </summary>
        private string WFProcessName;
        /// <summary>
        /// Project name param name in Custom parameters
        /// </summary>
        const string projectNameField = @"ProjectName";
        /// <summary>
        /// Default Project name to be used in case of missing project name
        /// </summary>
        const string defaultProjectName = @"Default";
        /// <summary>
        /// HPC Job name, field name in custom parameters table
        /// </summary>
        const string hpcJobFieldName = @"HpcJobName";
        /// <summary>
        /// HPC Task name, field name in custom parameters table
        /// </summary>
        const string hpcTaskFieldName = @"HpcTaskName";
        /// <summary>
        /// Domain for HPC specific custom parameters
        /// </summary>
        const string hpcDomainFieldName = @"HPC";
      
        /// <summary>
        /// HPC job id field name
        /// </summary>
        const string hpcJobIDFieldName = @"HpcJobID";

        /// <summary>
        /// HPC runat node field name
        /// </summary>
        const string hpcRunAtNode = @"RunATNode";

        /// <summary>
        /// HPC reserve node field name
        /// </summary>
        const string hpcReserveNode = @"ReserveNode";

        
        #endregion

        #region ITridentScheduler Members

        /// <summary>
        /// Initializes the Hpc Launcher from the App.Config file
        /// This funtion is expected to throw exception in case of error, 
        /// WfServiceScheduler will terminate the host process in case of exception
        /// </summary>
        public void Initialize(string executionProcName,
                                        ER.Connection registryConn,
                                        EventHandler exitEvent,
                                        EventHandler<WFLauchEventArgs> launchWorkflowUIHandler)
        {
                this.WFProcessName = executionProcName;
                registryConnection = registryConn;

                currentScheduleEngines = new List<ITridentHpcScheduleEngine>();
                //Try loading the schedule engine
                LoadEngine();

                ER.Machine headNode = ER.Machine.GetCreateForCluster(Environment.MachineName, registryConn);
                //In case of null path 
                if (String.IsNullOrEmpty(headNode.ComputeInstalledDirectory))
                {
                    headNode.ComputeInstalledDirectory = string.Empty;
                }
                // Initialize the cluster manager object with the head node, list of compute nodes and 
                // location of the workflow launcher exe (commandline job)
                clusterManager.Initialize(Environment.MachineName, Path.Combine(headNode.ComputeInstalledDirectory, this.WFProcessName));

                UpdateClusterNodesOnRegistry(headNode);

                clusterManager.OnTaskStatusChanged += new HPCSchedulerWrapper.TaskStatusChanged(OnTaskStatusChanged);
                clusterManager.OnJobStatusChanged += new HPCSchedulerWrapper.JobStatusChanged(OnJobStatusChanged);
        }

        /// <summary>
        /// Updates the nodes into registry
        /// </summary>
        private ER.Machine UpdateClusterNodesOnRegistry(ER.Machine headNode)
        {
            headNode.Type = ER.MachineType.HPCCluster;
            headNode.Parent = null;
            headNode.MaxConcurrentJobs = clusterManager.GetNumberOfCores(Environment.MachineName);
            headNode.Availability = ER.ExecutorAvailabiltiy.Online;
            headNode.Save();

            // Populate the compute nodes.
            foreach (string computeNode in clusterManager.ComputeNodeList)
            {
                ER.Machine node = null;
                try
                {
                    node = ER.Machine.Load(computeNode, registryConnection);
                    node.Parent = headNode;
                    node.MaxConcurrentJobs = clusterManager.GetNumberOfCores(computeNode);
                    node.Availability = ER.ExecutorAvailabiltiy.Online;
                    node.Save();
                }
                catch (ER.RecordNotFoundException)
                {
                    node = ER.Machine.Create(computeNode, ER.MachineType.HPCComputeNode, registryConnection);
                    node.Parent = headNode;
                    node.MaxConcurrentJobs = clusterManager.GetNumberOfCores(computeNode);
                    ER.HealthCheck.Create(node, registryConnection);
                    node.Availability = ER.ExecutorAvailabiltiy.Online;
                    node.Save();
                }
                
            }
            headNode.Refresh();
            return headNode;
        }
        /// <summary>
        /// Funtion called for processing set of jobs on cluster
        /// re-routes the jobs to the configured engine to get the schedule 
        /// Calls the local scheduleJob method to actual processing the job
        /// Excepted to throw exceptin which should be handled by the caller
        /// </summary>
        /// <param name="jobs"></param>
        public void ProcessJobs(bool showConsole, Collection<ER.Job> jobs, ER.Machine machine)
        {
            //Connect to registry, if not connected already
            ReOpenRegistryConnection();
            //Check for update from management studio system, will discard the changes happening in between
            string runtimePath = Path.Combine(machine.ComputeInstalledDirectory, this.WFProcessName);
            if (!clusterManager.TridentLauncherPath.Equals(runtimePath, StringComparison.CurrentCultureIgnoreCase))
            {
                clusterManager.TridentLauncherPath = runtimePath;
            }

            foreach (ER.Job tridentJob in jobs)
            {
                bool jobProcessed = false;
                //Do a refresh to get any changes to machine config
                machine.Refresh();
                //Check whether the user name and password is configured
                if (string.IsNullOrEmpty(machine.ClusterUser) || string.IsNullOrEmpty(machine.ClusterPassword))
                {
                    tridentJob.ErrorMessage = string.Format("HPC Cluster User Name or Password is not configured, please use Trident Management studio -> Registry manager -> Nodes to configure valid user for Cluster Node : {0}", machine.Name);
                    tridentJob.Status = ER.JobStatus.Aborted;
                    continue;
                }
                if (tridentJob.Machines[0].Type == ER.MachineType.HPCCluster)
                {
                    //Engine search required only if the specified node is a cluster
                    foreach (ITridentHpcScheduleEngine engine in currentScheduleEngines)
                    {
                        try
                        {
                            string[] runAtNodes = engine.FindNodes(tridentJob, tridentJob.Machines[0], clusterManager.ComputeNodeList.ToArray());
                            if (runAtNodes.Length >= 1)
                            {
                                jobProcessed = true;
                                //Schedule the workflow execution to the node provided by the schedule engine
                                ScheduleJob(tridentJob, runAtNodes);

                                //Call the post schedule interface
                                engine.PostScheduleJob(tridentJob, tridentJob.Machines[0], runAtNodes);
                                break;
                            }
                        }
                        //we should  not stop processing, if some engine throws any exception
                        catch (Exception exp)
                        {
                            Log("Error", exp.Message);
                        }
                    }
                    if (!jobProcessed)
                    {
                        jobProcessed = true;
                        //Schedule the workflow execution to the node provided by the schedule engine
                        ScheduleJob(tridentJob, clusterManager.ComputeNodeList.ToArray());
                        
                        break;
                    }

                }
                else if (tridentJob.Machines[0].Type == ER.MachineType.HPCComputeNode)
                {
                    jobProcessed = true;

                    string[] runAtNodes = new string[1];
                    runAtNodes[0] = tridentJob.Machines[0].Name;
                    //Schedule the workflow execution to the node provided by the job
                    ScheduleJob(tridentJob, runAtNodes);
                }                
            }
        }

        #endregion
        /// <summary>
        /// Event handler for job status notification from HPC scheduler
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="args"></param>
        void OnJobStatusChanged(object sender, SH.JobStateEventArg args)
        {
            if (SHP.JobState.Canceled == args.NewState ||
                SHP.JobState.Failed == args.NewState)
            {
                try
                {
                    ProcessFailure(sender, args.JobId, null);
                }
                catch (Exception exp)
                {
                    this.Log("Error", exp.Message);
                    this.Log("Error", exp.StackTrace);
                }
            }
 
        }

        /// <summary>
        /// Event handler for Task status notification from HPC scheduler
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="args"></param>
        void OnTaskStatusChanged(object sender, SH.TaskStateEventArg args)
        {
            if (SHP.TaskState.Canceled == args.NewState ||
                SHP.TaskState.Failed == args.NewState)
            {
                try
                {
                    ProcessFailure(sender, args.JobId, args.TaskId); 
                }
                catch (Exception exp)
                {
                    this.Log("Error", exp.Message);
                    this.Log("Error", exp.StackTrace);
                }
            }            
        }

        /// <summary>
        /// Process the failure notification from HPC to mark the trident job status as Aborted 
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="jobID"></param>
        /// <param name="taskID"></param>
        private void ProcessFailure(object sender, int jobID, SHP.ITaskId taskID)
        {
            SH.ISchedulerJob hpcJob = null;
            SH.ISchedulerTask hpcTask = null;
            SH.IScheduler scheduler = (SH.IScheduler)sender;
            //Open the job from hpc to get the trident job id
            hpcJob = scheduler.OpenJob(jobID);
            SH.INameValueCollection propertyCollection = hpcJob.GetCustomProperties();
            //loop through the property collection 
            foreach (SH.INameValue property in propertyCollection)
            {
                //Find the trident job id custom property
                if (property.Name == HPCSchedulerWrapper.tridentJobIDFieldName)
                {
                    string tridentJobID = property.Value;

                    if (!String.IsNullOrEmpty(tridentJobID))
                    {
                        ER.Connection regConnection = null;
                        if (null != registryConnection && registryConnection.Alive)
                        {
                            //Clone it for use it in this thread.
                            regConnection = registryConnection.Clone();
                            regConnection.Open();
                        }
                        ER.Job tridentJob = ER.Job.Load(new Guid(tridentJobID), regConnection);
                        //Update the status as aborted
                        if (tridentJob.Status != ER.JobStatus.Aborted)
                        {
                            if (null != taskID)
                            {
                                hpcTask = hpcJob.OpenTask(taskID);
                                if (string.IsNullOrEmpty(tridentJob.ErrorMessage)
                                    && !string.IsNullOrEmpty(hpcTask.ErrorMessage))
                                {
                                    tridentJob.ErrorMessage = hpcTask.ErrorMessage;
                                }
                            }
                            tridentJob.Status = ER.JobStatus.Aborted;
                        }
                        if (null != regConnection)
                        {
                            regConnection.Close();
                        }
                    }
                    break;
                }
            }
        }
        /// <summary>
        /// 
        /// </summary>
        private void LoadEngine()
        {
            DirectoryInfo dirInfo = new DirectoryInfo(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location));
            //Enumerate the engine dlls available in the system
            FileInfo[] fileinfos = dirInfo.GetFiles("*.dll");
            //Add all the engines to the list
            foreach (FileInfo fInfo in fileinfos)
            {
                try
                {
                    Assembly asm = Assembly.LoadFile(fInfo.FullName);
                    Type[] asmTypes = asm.GetTypes();
                    foreach (Type t in asmTypes)
                    {
                        Type[] ifaces = t.GetInterfaces();
                        if (null != ifaces && ifaces.Contains(typeof(ITridentHpcScheduleEngine)))
                        {
                            //Found the class which has Schedule engine implementation, create an object and keep it
                            ITridentHpcScheduleEngine engine = (ITridentHpcScheduleEngine)Activator.CreateInstance(t);
                            currentScheduleEngines.Add(engine);
                            break;
                        }
                    }
                }
                catch(Exception)
                {
                }
            }            
        }

        /// <summary>
        /// This function is used to check and 
        /// reopen the registry connection if it is closed.
        /// </summary>
        private void ReOpenRegistryConnection()
        {
            if (null == this.registryConnection || !this.registryConnection.Alive)
            {
                ER.ConnectionManager mgr = ER.ConnectionManager.CreateForAgent(
                    ER.ConnectionManager.CancelBehavior.ThrowException,
                    ER.ConnectionManager.Agent.ExecutionService);

                this.registryConnection = mgr.PickConnection(ER.ConnectionManager.ConnectionUI.NeverShowUI);
            }
        }

        /// <summary>
        /// This method will be invoked to schedule the job into the HPC scheduler
        /// In case of any exception the job will be marked as Aborted with the error message
        /// </summary>
        /// <param name="tridentJob"></param>
        /// <param name="selectedNode"></param>
        void ScheduleJob(ER.Job tridentJob, string[] selectedNodes)
        {
            try
            {
                string workflowName = tridentJob.ActivityInstances[0].Activity.Name;
                string projectName = defaultProjectName;
                //Find the project name from the custom parameters
                foreach (ER.AdditionalParameter custParam in tridentJob.AdditionalParameters)
                {
                    if (0 == string.Compare(custParam.ParamName, projectNameField, true))
                    {
                        projectName = custParam.ParamValue;
                        break;
                    }
                }
                
                SHP.JobPriority priority = SHP.JobPriority.Normal;
                switch (tridentJob.Priority)
                {
                    case ER.JobPriority.High:
                        priority = SHP.JobPriority.AboveNormal;
                        break;
                    case ER.JobPriority.Medium:
                        priority = SHP.JobPriority.Normal;
                        break;
                    case ER.JobPriority.Low:
                        priority = SHP.JobPriority.BelowNormal;
                        break;
                    default:
                        priority = SHP.JobPriority.Normal;
                        break;
                 }

                //Create the Job/Task name for HPC
                string jobName = string.Format("Trident:{0}:{1}:{2}:TRI-{3}", projectName, workflowName, selectedNodes[0], ((ER.IObject)tridentJob).ID.ToString());
                if (jobName.Length > 78)
                {
                    jobName = jobName.Substring(0, 78);
                }

                int concurrentJobs = tridentJob.Machines[0].MaxConcurrentJobs;
                //Create the job
                SH.ISchedulerJob hpcJob = clusterManager.CreateJob( selectedNodes,
                                                    ((ER.IObject)tridentJob).ID.ToString(),
                                                    projectName,
                                                    jobName,
                                                    priority,
                                                    tridentJob.Exclusive,
                                                    concurrentJobs);
                
                clusterManager.CreateTask(((ER.IObject)tridentJob).ID.ToString(),
                                            selectedNodes,
                                            jobName,
                                            hpcJob);

                clusterManager.Scheduler.AddJob(hpcJob);

                string headNodeUser;
                string headNodePass;

                GetClusterUserInfo(tridentJob, out headNodeUser, out headNodePass);

                //Submit the job using the headnode user and password
                clusterManager.Scheduler.SubmitJob(hpcJob,
                                                headNodeUser,
                                                headNodePass);

                CreateJobCustomParam(tridentJob, hpcDomainFieldName, hpcJobIDFieldName, hpcJob.Id.ToString());

                CreateJobCustomParam(tridentJob, hpcDomainFieldName, hpcJobFieldName, jobName);

                CreateJobCustomParam(tridentJob, hpcDomainFieldName, hpcTaskFieldName, jobName);
                
                //Status changes happened between the time we picked the job and reaching here will be ignored
                //If the job is aborted this will be updated with the new status update logic
                tridentJob.Status = ER.JobStatus.StartPending;
                tridentJob.Save();
            }
            catch (Exception exp)
            {
                string errorMessage = string.Format("Scheduler thrown an exception " +
                            "while processing Job : {0} with Jobid:{1} : details: {2}", 
                            tridentJob.Name,
                            ((ER.IObject)tridentJob).ID.ToString(),
                            exp.Message + exp.StackTrace.ToString());
                try
                {
                    tridentJob.Status = ER.JobStatus.Aborted;
                    tridentJob.ErrorMessage = errorMessage;
                    tridentJob.Save();
                }
                catch (Exception ex)
                {
                    errorMessage += ex.Message;
                }
                this.Log("Error", errorMessage);
            }
        }

        #region ILoggable Members

        /// <summary>
        /// Fired when a log message is received.
        /// </summary>
        public event EventHandler<LogEventArgs> LogMessageReceived;

        /// <summary>
        /// Raises the LogMessageReceived event.
        /// </summary>
        /// <param name="level">The log level.</param>
        /// <param name="message">The log message.</param>
        protected virtual void OnLogMessageReceived(string level, string message)
        {
            this.LogMessageReceived.Fire(this, new LogEventArgs(level, message));
        }

        /// <summary>
        /// Writes the log message.
        /// </summary>
        /// <param name="message">Message string to be logged.</param>
        private void Log(string level, string message)
        {
            this.OnLogMessageReceived(level, String.Format(CultureInfo.CurrentUICulture, "ServerScheduler: {0}", message));
        }

        /// <summary>
        /// Writes the log message.
        /// </summary>
        /// <param name="format">Format of the message.</param>
        /// <param name="args">Arguments to be formatted.</param>
        private void Log(string level, string format, params object[] args)
        {
            this.Log(level, String.Format(CultureInfo.CurrentUICulture, format, args));
        }
        #endregion

        /// <summary>
        /// Gets the user information from registry for the specified cluster head node
        /// </summary>
        /// <param name="headUser"></param>
        /// <param name="headPass"></param>
        void GetClusterUserInfo(ER.Job tridentJob, out string headUser, out string headPass)
        {
            ER.Machine node = tridentJob.Machines[0];
            ER.Machine headNode = node;
            //find the head node if the selected node is a compute node
            if (node.Type == ER.MachineType.HPCComputeNode)
            {
                headNode = node.Parent;
            }
            headUser = headNode.ClusterUser;
            headPass = headNode.ClusterPassword;
        }
        /// <summary>
        /// Creates the custom parameter for the job in registry with details provided
        /// </summary>
        /// <param name="domain"></param>
        /// <param name="paramName"></param>
        /// <param name="paramValue"></param>
        void CreateJobCustomParam(ER.Job tridentJob, string domain, string paramName, string paramValue)
        {
            //Update HPC job id to custom Param
            ER.AdditionalParameter jobParam = ER.AdditionalParameter.Create(tridentJob, this.registryConnection);
            jobParam.Domain = domain;
            jobParam.ParamName = paramName;
            jobParam.ParamValue =paramValue;
            jobParam.Save(); 
        }

        #region IScheduleStatusUpdate members
        /// <summary>
        /// Gets the status of the job and tasks from HPC
        /// </summary>
        /// <param name="job"></param>
        /// <returns></returns>
        public ER.JobStatus GetScheduleStatus(ER.Job job)
        {
            string hpcJobID = null;
            string hpcTaskName = null;
            ER.JobStatus status = ER.JobStatus.Aborted;
            //Get the hpcJobId and taskid
            foreach (ER.AdditionalParameter param in job.AdditionalParameters)
            {
                if (param.Domain.Equals(hpcDomainFieldName, StringComparison.InvariantCultureIgnoreCase))
                {
                    if (param.ParamName.Equals(hpcJobIDFieldName, StringComparison.InvariantCultureIgnoreCase))
                    {
                        hpcJobID = param.ParamValue;                        
                    }
                    if (param.ParamName.Equals(hpcTaskFieldName, StringComparison.InvariantCultureIgnoreCase))
                    {
                        hpcTaskName = param.ParamValue;
                    }
                    //Get out of loop when we get the required two parameters
                    if (null != hpcJobID && hpcTaskName != null)
                    {
                        break;
                    }
                }             
            }
            if (!string.IsNullOrEmpty(hpcJobID) && !string.IsNullOrEmpty(hpcTaskName))
            {
                try
                {
                    //Get the actual HPC job status
                    SHP.TaskState taskState = clusterManager.GetTaskState( Convert.ToInt32(hpcJobID), hpcTaskName);
                    
                    //This request will be coming only when the job status is marked as running in trident system
                    //If HPC status is Finished or failed or Cancelled, 
                    //Mark the job stauts as aborted, only after checking the current status with registry once again
                    status = (taskState == SHP.TaskState.Failed || 
                            taskState == SHP.TaskState.Finished ||
                            taskState == SHP.TaskState.Canceled) ? ER.JobStatus.Aborted : job.Status;
                }
                catch (Exception exp)
                {
                    this.Log("Error", exp.Message);
                }
            }
            return status;
        }

        #endregion


        #region ITridentScheduler Members


        public void UpdateRegistryConnection(ER.Connection registryConn)
        {
            this.registryConnection = registryConn;
        }

        #endregion
    }
}
